# -*- coding: utf-8 -*-
"""
@project: 图片切割 for zuo
@author: Wu Yue
@file: reSplit
@ide: PyCharm
@creat time: 2020-08-20 10:07
@change time:
@function: 将 512 × 512 的原始云图及其label 重新分割为 256 × 256 的图片  图片与label需要保持对应关系
"""

import os, re, cv2
import numpy as np
import matplotlib.pyplot as plt
import shutil


""" 参数设置 """
# >>> path >>>
rootDir = "./outputPic"
sourceDirs = ["/grayPics"]
picSubDir = "/data"
labelSubDir = "/Label"
outputDir = "./reSplit"
# <<< path <<<

# >>> reguExp >>>
picRegu = "wind[0-9]*_[0-9]*_[0-9]*.jpg"
labelRegu = "wind[0-9]*_[0-9]*_[0-9]*.png"
# <<< reguExp <<<

# >>> file format >>>
picFormat = "jpg"
labelFormat = "png"
# <<< file format <<<
""" 参数设置 """


for sourceDir in sourceDirs:
    for root, dirs, files in os.walk(rootDir + sourceDir + picSubDir):
        for file in files:
            if re.match(picRegu, file):
                picPath = root + "/" + file
                labelPath = rootDir + sourceDir + labelSubDir + "/" + file.split(".")[0] + "." + labelFormat
                if os.path.exists(labelPath):
                    pic = np.array(plt.imread(picPath))
                    shutil.copyfile(labelPath, "./tempLabel.jpg")
                    label = np.array(plt.imread("./tempLabel.jpg"))
                    if len(label.shape) == 3:  # RGBA
                        r = label[:, :, 0]
                        g = label[:, :, 1]
                        b = label[:, :, 2]
                        a = label[:, :, 3]

                        assert np.equal(r, g).all() & np.equal(r, b).all()
                        label = r

                    # 断言 图片尺寸都是512 * 512
                    assert (pic.shape == (512, 512)) & (label.shape == (512, 512))

                    picSplited = [pic[0:256, 0:256], pic[0:256, 256:], pic[256:, 0:256], pic[256:, 256:]]
                    labelSplited = [label[0:256, 0:256], label[0:256, 256:], label[256:, 0:256], label[256:, 256:]]

                    print("正在处理: {}".format(picPath))
                    for i in range(len(picSplited)):
                        newPicName = file.split(".")[0] + "_" + str(i+1) + "." + picFormat
                        newLabelName = file.split(".")[0] + "_" + str(i+1) + "." + labelFormat

                        newPicDir = outputDir + sourceDir + picSubDir
                        newLabelDir = outputDir + sourceDir + labelSubDir

                        if not os.path.exists(newPicDir):
                            os.makedirs(newPicDir)
                        if not os.path.exists(newLabelDir):
                            os.makedirs(newLabelDir)

                        cv2.imwrite(newPicDir + "/" + newPicName, picSplited[i])
                        cv2.imwrite(newLabelDir + "/" + newLabelName, labelSplited[i])
os.remove("./tempLabel.jpg")